import datetime
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.widgets import MultiCursor
from pandas.plotting import register_matplotlib_converters

import QUANTAXIS as QA
from QAStrategy.QAStrategy.qastockbase import QAStrategyStockBase


class Beili(QAStrategyStockBase):
    '''
    本策略判断
    '''
    def __init__(self, username, password, code='510300', start='2019-01-01', end='2019-10-21', show_trade=True,
                 N=120, deadband=0.1, min_trade_money=5000,
                 dynamic=True, weight_pos=None, weight_distance=None, err='std', atr_ratio=3,
                 frequence='1min', strategy_id='QA_STRATEGY', init_cash=1000000, market=QA.MARKET_TYPE.STOCK_CN):
        super().__init__(username=username, password=password, code=code, start=start, end=end,
                         frequence=frequence, strategy_id=strategy_id, init_cash=init_cash)
        if weight_distance is None:
            self.weight_distance = [-2, -1, 1, 2]
        else:
            self.weight_distance = weight_distance
        if weight_pos is None:
            # 设置网格的仓位
            # self.context_weight = [0.9, 0.7, 0.5, 0.3, 0.1]
            self.context_weight = [1.0, 0.8, 0.5, 0.2, 0.0]
        else:
            self.context_weight = weight_pos

        # print(self.weight_distance)
        # print(self.context_weight)
        self.market_type = market
        self.N = N
        self.deadband_ratio = deadband
        self.last_grid = -1
        self.dynamic = dynamic
        self.err = err
        self.atr_ratio = atr_ratio  # N数字越小，atr比例应该越低
        self.end = datetime.datetime.strptime(end, '%Y-%m-%d')
        self.trade_start_time = None
        self.show_Trade_log = show_trade
        self.min_trade_money = min_trade_money
        self.min2day_coef = 1
        pre_end = datetime.datetime.strptime(start, '%Y-%m-%d')
        if frequence == QA.FREQUENCE.DAY:
            # pre_start = pre_end + datetime.timedelta(days=-N*2)
            pre_start = QA.QA_util_get_pre_trade_date(pre_end, n=N+1)
        else:
            if frequence is QA.FREQUENCE.SIXTY_MIN:
                self.min2day_coef = 4
            elif frequence is QA.FREQUENCE.THIRTY_MIN:
                self.min2day_coef = 8
            elif frequence is QA.FREQUENCE.FIFTEEN_MIN:
                self.min2day_coef = 16
            elif frequence is QA.FREQUENCE.FIVE_MIN:
                self.min2day_coef = 48
            else:
                self.min2day_coef = 240
            pre_n = int(N / self.min2day_coef) + 1
            pre_start = QA.QA_util_get_pre_trade_date(pre_end, n=pre_n)
        self.start = datetime.datetime.strptime(pre_start, '%Y-%m-%d')
        self.pre_data = QA.QA_quotation(code, pre_start, pre_end.strftime('%Y-%m-%d'),
                                        frequence=frequence, market=self.market_type,
                                        source=QA.DATASOURCE.MONGO, output=QA.OUTPUT_FORMAT.DATASTRUCT)
        self.all_data = QA.QA_quotation(code, pre_start, end,
                                        frequence=frequence, market=self.market_type,
                                        source=QA.DATASOURCE.MONGO, output=QA.OUTPUT_FORMAT.DATASTRUCT)
        self.context_band_dynamic = self.all_data.add_func(QA.QA_indicator_GRID, N=N * self.min2day_coef)
        print(self.context_band_dynamic.iloc[-1])
        # 获取网格区间分界线
        self.context_band = np.mean(self.pre_data.data['close']) \
                            + np.array([-40, -3, -2, 2, 3, 40]) * np.std(self.pre_data.data['close'])
        self.context_band = [-2.8616755, 4.5, 4.7, 4.9, 5.0, 5.1, 5.2, 9.5001755]
        if not self.dynamic:
            print('当前网格区间为：',  self.context_band)

    # sim account on bar
    def on_bar(self, bars):
        current_time = datetime.datetime.now()
        bar = bars.iloc[0]
        pos = self.acc.get_position(self.code)
        self.last_grid = self.acc.pre_grid
        # simu时，不应该动态获取，应该只在初始化时，也就是每天早上生成一次网格
        self.context_band_dynamic = self.market_data.groupby(level=1, sort=False).apply(QA.QA_indicator_GRID, N=self.N*self.min2day_coef)
        # print(self.context_band_dynamic.iloc[-1])
        if self.dynamic:
            grid = self.get_dynamic_grid(bar, weight_line=self.weight_distance)
        else:
            # grid = pd.cut([bar.close], self.context_band, labels=[0, 1, 2, 3, 4])[0]
            grid = self.get_static_grid(bar)
        self.last_grid = grid
        amount = self.calc_vol_by_pos(bar.close, pos, self.acc.available, self.context_weight[grid])
        if self.enable_trade(amount * bar.close):
            if amount > 0:
                self.acc.pre_grid = grid
                self.send_order('BUY', 'OPEN', price=bar.close, volume=amount, code=self.code)
            elif amount < 0:
                self.acc.pre_grid = grid
                # sell must be close, otherwise position will sell twice times macong 2020.10.13
                self.send_order('SELL', 'CLOSE', price=bar.close, volume=-amount, code=self.code)

        print(current_time, self.code, pos.volume_long, pos.last_price, bar.close)
        self.acc.on_bar(bar)  # update the realtime price

    def on_bar_backtest(self, bar):
        bar['datetime'] = bar.name[0].strftime("%Y-%m-%d %H:%M:%S")
        bar['code'] = bar.name[1]
        pos = self.acc.get_position(self.code)
        self.last_grid = self.acc.pre_grid
        if self.dynamic:
            grid = self.get_dynamic_grid(bar, weight_line=self.weight_distance)
        else:
            # grid = pd.cut([bar.close], self.context_band, labels=[0, 1, 2, 3, 4])[0]
            grid = self.get_static_grid(bar)
        self.last_grid = grid
        # hold_vol = self.acc.hold.get(self.code, 0)
        hold_vol = pos.volume_long
        # sell all at end time
        if bar.name[0].day == self.end.day and bar.name[0].month == self.end.month and bar.name[0].year == self.end.year:
            if self.frequence is QA.FREQUENCE.DAY:
                self.sell(bar, hold_vol)
            else:
                if bar.name[0].hour == 15:  # min bar sell at 15:00
                    self.sell(bar, hold_vol)
        else:
            amount = self.calc_vol(bar.close, hold_vol, self.acc.cash_available, self.context_weight[grid], towards=0)
            # amount = self.calc_vol_by_pos(bar.close, pos, self.acc.available, self.context_weight[grid])
            if self.enable_trade(amount * bar.close):
                if self.trade_start_time is None:
                    self.trade_start_time = bar.name[0].strftime("%Y-%m-%d %H:%M:%S")
                if amount > 0:
                    self.buy(bar, amount, grid)
                elif amount < 0:
                    if -amount <= self.acc.sell_available.get(self.code, 0):  # min bar 当日没settle
                        self.sell(bar, -amount, grid)

    def calc_vol_by_pos(self, price, pos, cash, ratio):
        hold = price * (pos.volume_long)  # position money
        total = hold + cash
        hold_ratio = hold / total
        if ratio == 0:
            return -pos.volume_long_his
        elif hold_ratio > ratio:
            amount = int((hold - total * ratio) / price / 100) * 100
            if amount <= pos.volume_long_his:  # 卖出时需要考虑历史long仓位
                return -amount
            else:
                return 0
        elif hold_ratio < ratio:
            amount = int((total * ratio - hold) / price / 100.25) * 100
            return amount
        else:
            return 0

    def calc_vol(self, price, hold_vol, cash, ratio, towards=1):
        hold = price * hold_vol
        total = hold + cash
        hold_ratio = hold / total
        if hold_ratio != ratio:
            if towards != -1:  # BUY OR keep position grid
                return int((total * ratio - hold) / price / 100.25) * 100  # 100.25的意思是包含手续费
            else:
                return int((hold - total * ratio) / price / 100) * 100
        else:
            return 0

    def get_dynamic_grid(self, bar, weight_line=None):
        if weight_line is None:
            weight_line = [-3, -2, 2, 3]
        # dynamic_value = self.context_band_dynamic.iloc[-1]
        dynamic_value = self.context_band_dynamic.loc[(bar.name[0], bar.name[1]), :]
        # print(dynamic_value[['MID_C', 'ATR']])
        if isinstance(dynamic_value, pd.DataFrame):
            dynamic_value = dynamic_value.iloc[0]
        std = 0
        if self.err == 'std':
            std = dynamic_value.STD_C
        elif self.err == 'atr':
            std = dynamic_value.ATR * self.atr_ratio  # atr的数值要比std小很多，这里给3倍参数
        mid = dynamic_value.MID_C  # 分钟bar中的VWAP要多除了100，而日bar则没这问题

        lines = [x * std + mid for x in weight_line]
        lines.append(1e5)
        lines.insert(0, 0)
        for index, value in enumerate(lines):
            if index == len(lines) - 1:
                break
            value_next = lines[index + 1]
            if value < bar.close <= value_next:
                if self.last_grid == index or self.last_grid == -1:  # 与上次在同一个格内或开仓
                    return index
                elif self.last_grid < index:  # 上次在下方格中
                    if bar.close > value + self.deadband_ratio * std:
                        return index
                    else:
                        return index - 1  # 这里不用last_grid是允许出现跳格情况
                else:  # 上次在上方格中
                    if bar.close < value_next - self.deadband_ratio * std:
                        return index
                    else:
                        return index + 1
        return self.last_grid

    def get_static_grid(self, bar, std=0.2):
        for index, value in enumerate(self.context_band):
            if index == len(self.context_band) - 1:
                break
            value_next = self.context_band[index + 1]
            if value < bar.close <= value_next:
                if self.last_grid == index or self.last_grid == -1:  # 与上次在同一个格内或开仓
                    return index
                elif self.last_grid < index:  # 上次在下方格中
                    if bar.close > value + self.deadband_ratio * std:
                        return index
                    else:
                        return index - 1  # 这里不用last_grid是允许出现跳格情况
                else:  # 上次在上方格中
                    if bar.close < value_next - self.deadband_ratio * std:
                        return index
                    else:
                        return index + 1
        return self.last_grid

    def plot_grid(self):
        data = pd.DataFrame()
        std = 0
        if self.err == 'std':
            std = self.context_band_dynamic.STD_C
        elif self.err == 'atr':
            std = self.context_band_dynamic.ATR * self.atr_ratio
        mid = self.context_band_dynamic.MID_C
        data['close'] = self.context_band_dynamic.CLOSE
        data['MID'] = mid

        if self.dynamic:
            for i in self.weight_distance:
                data['G{}'.format(i)] = mid + i * std
        else:
            data['U3_STATIC'] = self.context_band[6]
            data['U2_STATIC'] = self.context_band[5]
            data['U1_STATIC'] = self.context_band[4]
            data['L1_STATIC'] = self.context_band[3]
            data['L2_STATIC'] = self.context_band[2]
            data['L3_STATIC'] = self.context_band[1]
        data.plot()
        # multi = MultiCursor(figure.canvas, (axes1, axes2, axes3, axes4), color='r', lw=1)
        plt.show()

    def enable_trade(self, money):
        if abs(money) >= self.min_trade_money:
            return True
        else:
            return False

    def buy(self, bar, amount, grid):
        order = self.acc.send_order(
            code=self.code,
            time=bar.name[0],
            amount=amount,
            towards=QA.ORDER_DIRECTION.BUY,
            price=bar.close,
            order_model=QA.ORDER_MODEL.CLOSE,
            amount_model=QA.AMOUNT_MODEL.BY_AMOUNT
        )
        if isinstance(order, bool):
            print('order is bool')
            return
        self.broker.receive_order(QA.QA_Event(order=order, market_data=bar))
        trade_mes = self.broker.query_orders(self.acc.account_cookie, 'filled')
        res = trade_mes.loc[order.account_cookie, order.realorder_id]
        order.trade(res.trade_id, res.trade_price, res.trade_amount, res.trade_time)
        if self.show_Trade_log:
            print(bar.name[0], 'BUY', self.context_weight[grid], bar.close, amount,
              int(self.acc.cash_available), int(self.acc.cash_available + self.acc.hold.get(self.code, 0) * bar.close))

    def sell(self, bar, amount, grid=None):
        order = self.acc.send_order(
            code=self.code,
            time=bar.name[0],
            amount=amount,
            towards=QA.ORDER_DIRECTION.SELL,
            price=bar.close,
            order_model=QA.ORDER_MODEL.MARKET,
            amount_model=QA.AMOUNT_MODEL.BY_AMOUNT
        )
        # print
        if isinstance(order, bool):
            print('order is bool')
            return
        self.broker.receive_order(QA.QA_Event(order=order, market_data=bar))
        trade_mes = self.broker.query_orders(self.acc.account_cookie, 'filled')
        res = trade_mes.loc[order.account_cookie, order.realorder_id]
        order.trade(res.trade_id, res.trade_price, res.trade_amount, res.trade_time)
        if self.show_Trade_log:
            if grid is None:
                print(bar.name[0], 'SEL', 0.0, bar.close, amount, int(self.acc.cash_available),
                      int(self.acc.cash_available + self.acc.hold.get(self.code, 0) * bar.close))
            else:
                print(bar.name[0], 'SEL', self.context_weight[grid], bar.close, amount, int(self.acc.cash_available),
                      int(self.acc.cash_available + self.acc.hold.get(self.code, 0) * bar.close))


